import tensorflow as tf
tf.app.flags.DEFINE_string(flag_name='job_name', default_value='ps', docstring="启动的服务器是ps还是worker")
tf.app.flags.DEFINE_integer(flag_name='task_index', default_value=0, docstring='指定是ps还是worker服务器中的哪一台服务器')

flg = tf.app.flags.FLAGS


def main(agrs):
    # 分布式的计算
    """
    通过分布式的计算, 完成一个基本的矩阵相乘
    :param agrs: 
    :return: 
    """

    # 1. 创建集群对象
    # 通过字典的方式来对急群众的参数服务器和工作服务器完成基本的描述的, 描述的方式是 ip:port
    # 由于没有GPU的原因, 只需要通过多台电脑的cpu来完成具体的计算
    # 让mac作为worker 服务器
    # 让unbutn作为参数服务器
    cluster = tf.train.ClusterSpec({'ps': ['192.168.31.22:2223'], 'worker': ['192.168.31.21:2222']})
    # 2. 指定是ps还是worker中的具体的哪一台服务器创建一个服务
    server = tf.train.Server(cluster,job_name=flg.job_name,task_index=flg.task_index)
    # 使用了最大训练步数的钩子对象 需要指定一个全局的步数计数器
    global_step = tf.contrib.framework.get_or_create_global_step()

    # 3. 指定具体的运行设备
    # 参数服务器只需要等待接受参数即可, 不需要指定运行的设备
    # 只有worker服务器才需要指定计算的设备
    if flg.job_name == 'ps':
        # 是参数服务器 ,等待接收参数即可
        server.join()
    else:
        # 启动的服务器工作服务器
        # 指定工作服务器中的哪个设备来进行计算, (计算的设备: CPU, GPU, TPU)
        worker_device = '/job:worker/task:%d/cpu:0' % flg.task_index
        print(worker_device,"+++++++")
        with tf.device(tf.train.replica_device_setter(worker_device=worker_device, cluster=cluster)):
            # 完成矩阵的相乘
            # 1行4列 * 4行一列 = 1行一列
            x = tf.Variable([[1,2,3,4]])
            y = tf.Variable([[3],[3],[3],[3]])
            mat_op = tf.matmul(x, y)

    # 4. 创建分布式的会话, 完成矩阵相乘的操作
    # 创建分布式的会话对象, 需要指定主worker
    with tf.train.MonitoredTrainingSession(
            master='grpc://192.168.31.21:2222',
            is_chief=True,
            config=tf.ConfigProto(log_device_placement=True),
            hooks=[tf.train.StopAtStepHook(last_step=5000)]
    ) as mon_sess:
        while not mon_sess.should_stop():
            # 使用 MonitoredTrainingSession 不需要创建 初始化的op对象
            ret = mon_sess.run(mat_op)
            print(ret)


if __name__ == '__main__':
    tf.app.run()